#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Neo4j决策树导入工具
功能：读取目录下的文件，根据->符号建立决策树节点关系，导入到Neo4j数据库
"""

import os
import re
import jieba
from pathlib import Path
from neo4j import GraphDatabase
from datetime import datetime
import hashlib


class DecisionTreeImporter:
    def __init__(self, uri="bolt://192.168.1.169:7687", username="neo4j", password="Pa@ssw0rd"):
        try:
            self.driver = GraphDatabase.driver(uri, auth=(username, password))
            print(f"✅ 成功连接到Neo4j数据库: {uri}")
        except Exception as e:
            print(f"❌ 连接Neo4j失败: {e}")
            self.driver = None
    
    def close(self):
        if self.driver:
            self.driver.close()
    
    def clear_database(self):
        """清空数据库"""
        if not self.driver:
            return False
        
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
            print("✅ 数据库已清空")
            return True
    
    def scan_directory(self, directory_path):
        """扫描目录下的所有txt文件"""
        directory_path = Path(directory_path)
        if not directory_path.exists():
            print(f"❌ 目录 {directory_path} 不存在")
            return []
        
        txt_files = list(directory_path.glob("*.txt"))
        print(f"📁 找到 {len(txt_files)} 个txt文件")
        
        all_lines = []
        for txt_file in txt_files:
            try:
                # 尝试UTF-8编码
                with open(txt_file, 'r', encoding='utf-8') as file:
                    lines = file.readlines()
                
                file_lines = []
                for line_num, line in enumerate(lines, 1):
                    line_content = line.strip()
                    if line_content and ' -> ' in line_content:  # 只处理包含决策树格式的行
                        file_lines.append({
                            'filename': txt_file.name,
                            'filepath': str(txt_file),
                            'line_number': line_num,
                            'content': line_content
                        })
                
                all_lines.extend(file_lines)
                print(f"✅ {txt_file.name}: {len(file_lines)} 行决策树数据")
                
            except UnicodeDecodeError:
                # 尝试GBK编码
                try:
                    with open(txt_file, 'r', encoding='gbk') as file:
                        lines = file.readlines()
                    
                    file_lines = []
                    for line_num, line in enumerate(lines, 1):
                        line_content = line.strip()
                        if line_content and ' -> ' in line_content:
                            file_lines.append({
                                'filename': txt_file.name,
                                'filepath': str(txt_file),
                                'line_number': line_num,
                                'content': line_content
                            })
                    
                    all_lines.extend(file_lines)
                    print(f"✅ {txt_file.name} (GBK): {len(file_lines)} 行决策树数据")
                    
                except Exception as e:
                    print(f"❌ 读取文件失败: {txt_file.name} - {e}")
            
            except Exception as e:
                print(f"❌ 读取文件失败: {txt_file.name} - {e}")
        
        print(f"📊 总计读取 {len(all_lines)} 行决策树数据")
        return all_lines
    
    def parse_decision_line(self, line_data):
        """解析决策树行，根据->建立节点"""
        content = line_data['content']
        
        # 分割决策路径
        parts = content.split(' -> ')
        
        if len(parts) < 2:
            return None
        
        # 提取产品编号（第一个部分）
        product_code = parts[0].strip()
        
        # 提取货号
        cargo_match = re.search(r'\[货号:([^\]]+)\]', content)
        cargo_number = cargo_match.group(1) if cargo_match else product_code
        
        # 清理最后一个部分的货号信息
        last_part = parts[-1]
        if cargo_match:
            last_part = last_part.replace(cargo_match.group(0), '').strip()
            parts[-1] = last_part
        
        # 构建节点链
        nodes = []
        for i, part in enumerate(parts):
            clean_part = part.strip()
            if clean_part:
                node_id = self._generate_node_id(clean_part, i)
                nodes.append({
                    'id': node_id,
                    'content': clean_part,
                    'level': i,
                    'type': self._classify_node_type(clean_part, i),
                    'keywords': ' '.join(jieba.cut(clean_part))
                })
        
        return {
            'product_code': product_code,
            'cargo_number': cargo_number,
            'nodes': nodes,
            'filename': line_data['filename'],
            'line_number': line_data['line_number'],
            'original_content': content
        }
    
    def _generate_node_id(self, content, level):
        """生成节点ID"""
        # 使用内容和级别生成唯一ID
        content_hash = hashlib.md5(content.encode('utf-8')).hexdigest()[:8]
        return f"node_{level}_{content_hash}"
    
    def _classify_node_type(self, content, level):
        """分类节点类型"""
        if level == 0:
            return 'product'
        elif '？' in content or '?' in content:
            return 'question'
        elif '建议' in content or '推荐' in content or '应该' in content:
            return 'solution'
        elif content in ['是', '否', 'Y', 'N', 'yes', 'no']:
            return 'answer'
        elif level == 1:
            return 'problem'
        elif level == 2:
            return 'category'
        elif '则' in content or '如果' in content:
            return 'condition'
        else:
            return 'step'
    
    def import_decision_trees(self, parsed_data_list):
        """导入决策树到Neo4j"""
        if not self.driver:
            print("❌ 数据库连接不可用")
            return False
        
        with self.driver.session() as session:
            imported_count = 0
            
            for data in parsed_data_list:
                try:
                    self._import_single_decision_tree(session, data)
                    imported_count += 1
                    
                    if imported_count % 50 == 0:
                        print(f"📊 已导入 {imported_count} 条决策树...")
                
                except Exception as e:
                    print(f"❌ 导入失败: {data['filename']}:{data['line_number']} - {e}")
            
            print(f"✅ 成功导入 {imported_count} 条决策树")
            return True
    
    def _import_single_decision_tree(self, session, data):
        """导入单个决策树"""
        # 1. 创建所有节点
        for node in data['nodes']:
            session.run("""
                MERGE (n:DecisionNode {id: $node_id})
                SET n.content = $content,
                    n.type = $type,
                    n.level = $level,
                    n.keywords = $keywords,
                    n.product_code = $product_code,
                    n.filename = $filename,
                    n.line_number = $line_number,
                    n.cargo_number = $cargo_number,
                    n.timestamp = $timestamp
            """,
            node_id=node['id'],
            content=node['content'],
            type=node['type'],
            level=node['level'],
            keywords=node['keywords'],
            product_code=data['product_code'],
            filename=data['filename'],
            line_number=data['line_number'],
            cargo_number=data['cargo_number'],
            timestamp=datetime.now().isoformat()
            )
        
        # 2. 创建节点之间的关系（根据->符号的顺序）
        for i in range(len(data['nodes']) - 1):
            current_node = data['nodes'][i]
            next_node = data['nodes'][i + 1]
            
            session.run("""
                MATCH (current:DecisionNode {id: $current_id})
                MATCH (next:DecisionNode {id: $next_id})
                MERGE (current)-[:LEADS_TO {
                    filename: $filename,
                    line_number: $line_number,
                    step: $step
                }]->(next)
            """,
            current_id=current_node['id'],
            next_id=next_node['id'],
            filename=data['filename'],
            line_number=data['line_number'],
            step=i + 1
            )
        
        # 3. 创建产品节点（如果不存在）
        session.run("""
            MERGE (p:Product {code: $product_code})
            SET p.cargo_number = $cargo_number,
                p.name = $product_code
        """, 
        product_code=data['product_code'],
        cargo_number=data['cargo_number']
        )
        
        # 4. 创建产品到第一个决策节点的关系
        if data['nodes']:
            first_node = data['nodes'][0]
            session.run("""
                MATCH (p:Product {code: $product_code})
                MATCH (n:DecisionNode {id: $node_id})
                MERGE (p)-[:HAS_DECISION_TREE {
                    filename: $filename,
                    line_number: $line_number
                }]->(n)
            """,
            product_code=data['product_code'],
            node_id=first_node['id'],
            filename=data['filename'],
            line_number=data['line_number']
            )
        
        # 5. 创建文件节点
        session.run("""
            MERGE (f:File {name: $filename})
            SET f.path = $filepath
        """, 
        filename=data['filename'],
        filepath=data.get('filepath', '')
        )
        
        # 6. 创建原始记录节点 - content是最后一个->右边的内容
        record_id = f"{data['filename']}_{data['line_number']}"

        # 获取最后一个节点的内容作为Record的content
        final_content = data['nodes'][-1]['content'] if data['nodes'] else ""

        session.run("""
            MERGE (r:Record {id: $record_id})
            SET r.content = $content,
                r.full_path = $full_path,
                r.filename = $filename,
                r.line_number = $line_number,
                r.product_code = $product_code,
                r.cargo_number = $cargo_number,
                r.timestamp = $timestamp
        """,
        record_id=record_id,
        content=final_content,
        full_path=data['original_content'],
        filename=data['filename'],
        line_number=data['line_number'],
        product_code=data['product_code'],
        cargo_number=data['cargo_number'],
        timestamp=datetime.now().isoformat()
        )
        
        # 7. 建立文件到记录的关系
        session.run("""
            MATCH (f:File {name: $filename})
            MATCH (r:Record {id: $record_id})
            MERGE (f)-[:CONTAINS]->(r)
        """,
        filename=data['filename'],
        record_id=record_id
        )
        
        # 8. 建立记录到决策节点的关系
        for node in data['nodes']:
            session.run("""
                MATCH (r:Record {id: $record_id})
                MATCH (n:DecisionNode {id: $node_id})
                MERGE (r)-[:INCLUDES]->(n)
            """,
            record_id=record_id,
            node_id=node['id']
            )
    
    def create_indexes(self):
        """创建索引以提高查询性能"""
        if not self.driver:
            return
        
        with self.driver.session() as session:
            indexes = [
                "CREATE INDEX product_code_idx IF NOT EXISTS FOR (p:Product) ON (p.code)",
                "CREATE INDEX decision_node_id_idx IF NOT EXISTS FOR (n:DecisionNode) ON (n.id)",
                "CREATE INDEX decision_node_content_idx IF NOT EXISTS FOR (n:DecisionNode) ON (n.content)",
                "CREATE INDEX decision_node_type_idx IF NOT EXISTS FOR (n:DecisionNode) ON (n.type)",
                "CREATE INDEX decision_node_product_idx IF NOT EXISTS FOR (n:DecisionNode) ON (n.product_code)",
                "CREATE INDEX file_name_idx IF NOT EXISTS FOR (f:File) ON (f.name)",
                "CREATE INDEX record_id_idx IF NOT EXISTS FOR (r:Record) ON (r.id)"
            ]
            
            for index_query in indexes:
                try:
                    session.run(index_query)
                    print(f"✅ 索引创建成功")
                except Exception as e:
                    print(f"⚠️  索引创建失败: {e}")
    
    def get_statistics(self):
        """获取导入统计信息"""
        if not self.driver:
            return
        
        with self.driver.session() as session:
            # 统计节点数量
            result = session.run("MATCH (n:DecisionNode) RETURN count(n) as node_count")
            node_count = result.single()['node_count']
            
            # 统计关系数量
            result = session.run("MATCH ()-[r:LEADS_TO]->() RETURN count(r) as relation_count")
            relation_count = result.single()['relation_count']
            
            # 统计产品数量
            result = session.run("MATCH (p:Product) RETURN count(p) as product_count")
            product_count = result.single()['product_count']
            
            # 统计文件数量
            result = session.run("MATCH (f:File) RETURN count(f) as file_count")
            file_count = result.single()['file_count']
            
            # 统计记录数量
            result = session.run("MATCH (r:Record) RETURN count(r) as record_count")
            record_count = result.single()['record_count']
            
            print(f"\n📊 导入统计信息:")
            print(f"   - 决策节点数: {node_count}")
            print(f"   - 决策关系数: {relation_count}")
            print(f"   - 产品数: {product_count}")
            print(f"   - 文件数: {file_count}")
            print(f"   - 记录数: {record_count}")


def main():
    """主函数"""
    import sys

    print("🚀 Neo4j决策树导入工具")
    print("=" * 50)

    # 配置参数
    directory_path = "/Users/daijunxiong/Downloads/knowledge_txt"
    auto_clear = len(sys.argv) > 1 and sys.argv[1] == "--auto-clear"

    print(f"📁 目录路径: {directory_path}")

    # 创建导入器
    importer = DecisionTreeImporter()

    if not importer.driver:
        print("❌ 无法连接到Neo4j数据库")
        return

    try:
        # 自动清空数据库或询问
        if auto_clear:
            print("🗑️  自动清空数据库...")
            importer.clear_database()
        else:
            choice = input("是否清空现有数据库？(y/N): ").strip().lower()
            if choice == 'y':
                importer.clear_database()
        
        # 扫描目录
        print(f"\n📁 扫描目录: {directory_path}")
        all_lines = importer.scan_directory(directory_path)
        
        if not all_lines:
            print("❌ 没有找到决策树数据")
            return
        
        # 解析决策树
        print("\n🔄 解析决策树数据...")
        parsed_data = []
        for line_data in all_lines:
            parsed = importer.parse_decision_line(line_data)
            if parsed:
                parsed_data.append(parsed)
        
        print(f"✅ 解析完成，共 {len(parsed_data)} 条决策树")
        
        # 导入到Neo4j
        print("\n💾 导入决策树到Neo4j...")
        success = importer.import_decision_trees(parsed_data)
        
        if success:
            # 创建索引
            print("\n🔧 创建索引...")
            importer.create_indexes()
            
            # 显示统计信息
            importer.get_statistics()
            
            print("\n🎉 导入完成！")
            print("\n💡 现在您可以使用Cypher查询决策树:")
            print("   MATCH (p:Product {code: 'PK511'})-[:HAS_DECISION_TREE*1..10]->(n)")
            print("   RETURN p, n")
            print("\n   MATCH path = (start:DecisionNode)-[:LEADS_TO*]->(end:DecisionNode)")
            print("   WHERE start.content CONTAINS '无产物'")
            print("   RETURN path")
        
    finally:
        importer.close()


if __name__ == "__main__":
    main()
